from .base import FactorBase
import baostock as bs
from common.variables import TimeFrameType
from datetime import datetime
import pandas as pd


class BaoKlineFactor(FactorBase):
    fname_list = [
        "open",
        "high",
        "low",
        "close",
        "volume",
        "amount",
        "adjustflag",
        "turn",
        "tradestatus",
        "pctChg",
        "peTTM",
        "pbMRQ",
        "psTTM",
        "pcfNcfTTM",
        "isST",
    ]

    def __init__(self, callbacks: list[callable] = None):
        self.callbacks = callbacks if callbacks else []

    @property
    def bfname(self):
        return self.basename

    def calculate(
        self,
        symbol: str,
        timeframe: TimeFrameType,
        start: str,
        end: str,
        callbacks: list[callable],
    ) -> pd.DataFrame:
        start_date = datetime.fromisoformat(start).strftime("%Y-%m-%d")
        end_date = datetime.fromisoformat(end).strftime("%Y-%m-%d")
        fields = (
            "date,code,open,high,low,close,volume,amount,"
            "adjustflag,turn,tradestatus,pctChg,peTTM,"
            "pbMRQ,psTTM,pcfNcfTTM,isST"
        )

        lg = bs.login()
        if lg.error_code != "0":
            print("login respond error_code:" + lg.error_code)
            print("login respond  error_msg:" + lg.error_msg)

        #### 获取沪深A股历史K线数据 ####
        # 详细指标参数，参见“历史行情指标参数”章节；“分钟线”参数与“日线”参数不同。“分钟线”不包含指数。
        # 分钟线指标：date,time,code,open,high,low,close,volume,amount,adjustflag
        # 周月线指标：date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg
        rs = bs.query_history_k_data_plus(
            symbol,
            fields,
            start_date=start_date,
            end_date=end_date,
            frequency="d",
            adjustflag="2",
        )
        if rs.error_code != "0":
            print("login respond error_code:" + rs.error_code)
            print("login respond  error_msg:" + rs.error_msg)

        data_list = []
        while (rs.error_code == "0") & rs.next():
            data_list.append(rs.get_row_data())
        result = pd.DataFrame(data_list, columns=rs.fields)
        bs.logout()
        df = result.rename({"date": "time"}).drop("code")
        for xfunc in self.callbacks:
            xfunc(symbol, df, timeframe)
        return df
